package com.star.spring.ratelimit.limiter;

import lombok.extern.slf4j.Slf4j;

import java.util.concurrent.atomic.AtomicLong;

/**
 * 自定义令牌桶实现
 * 
 * @author star
 * @since 2025-01-24
 */
@Slf4j
public class TokenBucket {
    
    private final long capacity;           // 桶容量
    private final long refillPeriod;       // 补充周期（秒）
    private final long tokensPerRefill;    // 每次补充的令牌数
    
    private final AtomicLong tokens;       // 当前令牌数
    private final AtomicLong lastRefillTime; // 上次补充时间
    
    /**
     * 构造函数
     * 
     * @param capacity 桶容量
     * @param maxRequests 最大请求数
     * @param timeWindowSeconds 时间窗口（秒）
     */
    public TokenBucket(long capacity, int maxRequests, int timeWindowSeconds) {
        this.capacity = capacity;
        this.refillPeriod = 1; // 每秒补充一次
        
        // 计算每秒补充的令牌数，使用浮点数避免整数除法问题
        double tokensPerSecond = (double) maxRequests / timeWindowSeconds;
        this.tokensPerRefill = Math.max(1, (long) Math.ceil(tokensPerSecond)); // 至少每秒补充1个令牌
        
        this.tokens = new AtomicLong(capacity); // 初始时桶是满的
        this.lastRefillTime = new AtomicLong(System.currentTimeMillis());
        
        log.info("TokenBucket created - capacity: {}, maxRequests: {}, timeWindow: {}s, tokensPerSecond: {}, tokensPerRefill: {}", 
                capacity, maxRequests, timeWindowSeconds, tokensPerSecond, tokensPerRefill);
    }
    
    /**
     * 尝试获取令牌
     * 
     * @param requestedTokens 请求的令牌数
     * @return 是否获取成功
     */
    public boolean tryAcquire(int requestedTokens) {
        refillTokens();
        
        long currentTokens = tokens.get();
        if (currentTokens >= requestedTokens) {
            long newTokens = currentTokens - requestedTokens;
            if (tokens.compareAndSet(currentTokens, newTokens)) {
                log.debug("Token acquired - requested: {}, remaining: {}", requestedTokens, newTokens);
                return true;
            }
        }
        
        log.debug("Token acquisition failed - requested: {}, available: {}", requestedTokens, currentTokens);
        return false;
    }
    
    /**
     * 补充令牌
     */
    private void refillTokens() {
        long now = System.currentTimeMillis();
        long lastRefill = lastRefillTime.get();
        long timePassed = now - lastRefill;
        
        if (timePassed >= refillPeriod) {
            // 计算需要补充的令牌数
            long refillCycles = timePassed / refillPeriod / 1000;
            long tokensToAdd = refillCycles * tokensPerRefill;
            
            if (tokensToAdd > 0) {
                long currentTokens = tokens.get();
                long newTokens = Math.min(capacity, currentTokens + tokensToAdd);
                
                if (tokens.compareAndSet(currentTokens, newTokens)) {
                    lastRefillTime.set(now);
                    log.debug("Tokens refilled - added: {}, total: {}", tokensToAdd, newTokens);
                }
            }
        }
    }
    
    /**
     * 获取当前令牌数
     */
    public long getCurrentTokens() {
        refillTokens();
        return tokens.get();
    }
    
    /**
     * 获取桶容量
     */
    public long getCapacity() {
        return capacity;
    }
    
    /**
     * 重置令牌桶
     */
    public void reset() {
        tokens.set(capacity);
        lastRefillTime.set(System.currentTimeMillis());
        log.info("TokenBucket reset - tokens: {}", capacity);
    }
}
